from typing import List, Union, Generator, Iterator
from openai import OpenAI

client = OpenAI(
    api_key="ragflow-k5MDM3NzA2MDk0ZjExZjBhOWRhMDI0Mm",
    base_url="http://172.16.3.123:8088/api/v1/chats_openai/a141a2003a0811f08d3a0242ac160006"
)


class Pipeline:
    def __init__(self):
        self.name: str = "My Chat"

    def pipe(
            self,
            user_message: str,
            model_id: str,
            messages: List[dict],
            body: Union[dict, list],
    ) -> Union[str, Generator, Iterator]:
        query = messages[-1]["content"]

        completion = client.chat.completions.create(
            model='',
            messages=messages,
            stream=True
        )

        for chunk in completion:
            content = chunk.choices[0].delta.content
            if content is not None:
                yield content
